package crypto

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
)

/**
 * RsaEncrypt，RSA 算法进行字符串加密
 * Params：
 *   plainText：待加密字符串的字节切片
 *   key：加密公钥字节切片
 * Return：
 *   加密结果字节切片
 */
func RsaEncrypt(plainText, publicKey []byte) ([]byte, error) {
	// pem解碼
	block, _ := pem.Decode(publicKey)
	pubInterface, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	// 斷言型別轉換
	pubKey := pubInterface.(*rsa.PublicKey)
	// 使用公鑰加密
	cipherText, err := rsa.EncryptPKCS1v15(rand.Reader, pubKey, plainText)
	if err != nil {
		return nil, err
	}
	return cipherText, nil
}

/**
 * RsaEncryptBase64，获取 RSA 加密后的 Base64 字符串
 * Params：
 *   plainText：待加密字符串的字节切片
 *   publicKey：RSA 公钥字节切片
 * Return：
 *   加密结果进行 Base64 编码后的字符串
 */
func RsaEncryptBase64(plainText, publicKey []byte) (string, error) {
	encrypted, err := RsaEncrypt(plainText, publicKey)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(encrypted), nil
}

/**
 * RsaDecrypt，RSA 解密
 * Params：
 *   cipherText：RSA 加密的数据字节切片
 *   key：加密秘钥
 * Return：
 *   解密后的明文字符串的字节切片
 */
func RsaDecrypt(cipherText, privateKey []byte) ([]byte, error) {
	// pem 解碼
	block, _ := pem.Decode(privateKey)
	privKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	// 使用私鑰解密
	plainText, err := rsa.DecryptPKCS1v15(rand.Reader, privKey, cipherText)
	if err != nil {
		return nil, err
	}
	return plainText, nil
}

/**
 * RsaDecryptBase64，根据 RSA 加密的 Base64 字符串获取解密的明文字符串切片
 * Params：
 *   cipherTextBase64：RSA 加密后进行 Base64 编码的字符串
 *   privateKey：RSA 私钥字节切片
 * Return：
 *   解密后的明文字符串字节切片
 */
func RsaDecryptBase64(cipherTextBase64 string, privateKey []byte) ([]byte, error) {
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return nil, err
	}
	return RsaDecrypt(cipherText, privateKey)
}

/**
 * GenerateRsaKey，生成 RSA 公私钥
 * Params：
 * Return：
 *   解密后的明文字符串字节切片
 */
func GenerateRsaKey() (prvKey, pubKey []byte) {
	// 生成私钥
	privateKey, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		panic(err)
	}
	derStream := x509.MarshalPKCS1PrivateKey(privateKey)
	block := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: derStream,
	}
	prvKey = pem.EncodeToMemory(block)
	// 生成公钥
	publicKey := &privateKey.PublicKey
	derPkix, err := x509.MarshalPKIXPublicKey(publicKey)
	if err != nil {
		panic(err)
	}
	block = &pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: derPkix,
	}
	pubKey = pem.EncodeToMemory(block)
	return
}

/**
 * RsaSignWithSha256，生成 RSA 签名信息
 * Params：
 *   plainText：待进行签名的明文字符串数据的字节切片
 *   priKey：RSA 私钥字节切片
 * Return：
 *   签名结果字符串字节切片
 */
func RsaSignWithSha256(plainText []byte, priKey []byte) ([]byte, error) {
	// 1. 获取 RSA 私钥信息
	block, _ := pem.Decode(priKey)
	if block == nil {
		return nil, errors.New("private key error")
	}
	privateKey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return nil, err
	}
	// 2. 数据进行 sha256 加密
	h := sha256.New()
	h.Write(plainText)
	hashed := h.Sum(nil)
	// 3. 生成签名信息
	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hashed)
	return signature, err
}

/**
 * RsaVerySignWithSha256，验证 RSA 签名
 * Params：
 *   plainText：待进行对比验证的明文字符串的字节切片
 *   signature：RSA 签名后的数据字节切片
 *   priKey：RSA 公钥字节切片
 * Return：
 *   签名结果字符串字节切片
 */
func RsaVerySignWithSha256(plainText, signature []byte, pubKey []byte) (bool, error) {
	// 1. 获取公钥信息
	block, _ := pem.Decode(pubKey)
	if block == nil {
		return false, errors.New("public key error")
	}
	publicKey, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return false, err
	}
	// 2. 对明文进行 hash256 运算得到散列值
	hashed := sha256.Sum256(plainText)
	// 3. 验签
	err = rsa.VerifyPKCS1v15(publicKey.(*rsa.PublicKey), crypto.SHA256, hashed[:], signature)
	if err != nil {
		return false, err
	}
	return true, nil
}
